import os
import json
import re
import hashlib
import numpy as np
from tqdm import tqdm
from transformers import AutoTokenizer
import multiprocessing
from functools import partial
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class ChineseCorpusPreprocessor:
    def __init__(self, config):
        """
        初始化预处理器
        :param config: 配置字典，包含所有处理参数
        """
        self.config = config
        self.tokenizer = AutoTokenizer.from_pretrained(
            config['tokenizer_name'],
            trust_remote_code=True,
            use_fast=config.get('use_fast_tokenizer', True)
        )
        self.seen_hashes = set()
        self.sensitive_words = self.load_sensitive_words(config.get('sensitive_words_path'))
        self.pii_patterns = self.load_pii_patterns()

    def load_sensitive_words(self, path):
        """加载敏感词表"""
        if not path or not os.path.exists(path):
            logger.warning(f"敏感词表文件不存在: {path}")
            return set()

        with open(path, 'r', encoding='utf-8') as f:
            return {line.strip().lower() for line in f if line.strip()}

    def load_pii_patterns(self):
        """加载PII(个人身份信息)检测正则表达式"""
        return [
            # 身份证号
            re.compile(r'[1-9]\d{5}(?:18|19|20)\d{2}(?:0[1-9]|1[0-2])(?:0[1-9]|[12]\d|3[01])\d{3}[\dXx]'),
            # 手机号
            re.compile(r'1[3-9]\d{9}'),
            # 银行卡号
            re.compile(r'\b[1-9]\d{9,18}\b'),
            # 邮箱
            re.compile(r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b'),
            # 中文地址特征
            re.compile(r'[\u4e00-\u9fa5]{2,10}(省|市|区|县|镇|乡|村|街道|路|号)')
        ]

    def clean_text(self, text):
        """
        文本清洗函数
        :param text: 原始文本
        :return: 清洗后的文本
        """
        # 1. 替换特殊空白字符
        text = re.sub(r'[\u200b-\u200d\u202a-\u202e]', ' ', text)

        # 2. 合并连续空白
        text = re.sub(r'\s+', ' ', text)

        # 3. 移除控制字符
        text = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', text)

        # 4. 简繁体统一 (实际应用时可用OpenCC)
        # text = opencc.convert(text)

        return text.strip()

    def should_keep(self, text):
        """
        判断文档是否应该保留
        :param text: 文本内容
        :return: bool, 是否保留
        """
        # 1. 长度检查
        char_len = len(text)
        if char_len < self.config['min_chars'] or char_len > self.config['max_chars']:
            logger.debug(f"长度不符: {char_len}")
            return False

        # 2. 中文比例检查
        chinese_chars = re.findall(r'[\u4e00-\u9fa5]', text)
        chinese_ratio = len(chinese_chars) / char_len if char_len > 0 else 0
        if chinese_ratio < self.config['min_chinese_ratio']:
            logger.debug(f"中文比例不足: {chinese_ratio:.2f}")
            return False

        # 3. 符号比例检查
        symbol_chars = re.findall(r'[^\w\s\u4e00-\u9fa5]', text)
        symbol_ratio = len(symbol_chars) / char_len if char_len > 0 else 0
        if symbol_ratio > self.config['max_symbol_ratio']:
            logger.debug(f"符号比例过高: {symbol_ratio:.2f}")
            return False

        # 4. 重复内容检查
        if self.config['max_repeat_ratio'] > 0:
            unique_chars = len(set(text))
            repeat_ratio = 1 - (unique_chars / char_len)
            if repeat_ratio > self.config['max_repeat_ratio']:
                logger.debug(f"重复内容过高: {repeat_ratio:.2f}")
                return False

        # 5. 敏感词检查
        text_lower = text.lower()
        if any(sw in text_lower for sw in self.sensitive_words):
            logger.debug("包含敏感词")
            return False

        # 6. PII检查
        if any(pattern.search(text) for pattern in self.pii_patterns):
            logger.debug("包含PII信息")
            return False

        # 7. 哈希去重
        text_hash = hashlib.md5(text.encode('utf-8')).hexdigest()
        if text_hash in self.seen_hashes:
            logger.debug("重复内容")
            return False

        self.seen_hashes.add(text_hash)
        return True

    def process_document(self, doc):
        """
        处理单个文档
        :param doc: 文档字典
        :return: 处理后的token IDs 或 None
        """
        try:
            text = doc.get('TEXT', '')
            if not text:
                return None

            # 清洗文本
            cleaned_text = self.clean_text(text)

            # 检查是否保留
            if not self.should_keep(cleaned_text):
                return None

            # Tokenization
            tokens = self.tokenizer.encode(
                cleaned_text,
                add_special_tokens=False,
                max_length=self.config['max_seq_length'],
                truncation=True
            )

            # 添加EOS token
            if self.tokenizer.eos_token_id is not None:
                tokens.append(self.tokenizer.eos_token_id)

            return tokens
        except Exception as e:
            logger.error(f"处理文档出错: {e}")
            return None

    def process_file(self, input_path, output_dir):
        """
        处理单个文件
        :param input_path: 输入文件路径
        :param output_dir: 输出目录
        """
        all_tokens = []
        output_path = os.path.join(output_dir, os.path.basename(input_path) + '.npy')

        logger.info(f"开始处理文件: {input_path}")

        # 读取并处理文档
        with open(input_path, 'r', encoding='utf-8') as f:
            for line in tqdm(f, desc=f"处理 {os.path.basename(input_path)}"):
                try:
                    doc = json.loads(line)
                    tokens = self.process_document(doc)
                    if tokens:
                        all_tokens.extend(tokens)
                except json.JSONDecodeError:
                    logger.warning(f"JSON解析失败: {line[:100]}...")

        # 保存为npy文件
        if all_tokens:
            token_array = np.array(all_tokens, dtype=np.int32)
            np.save(output_path, token_array)
            logger.info(f"已保存 {len(token_array)} tokens 到 {output_path}")
            return len(token_array)
        return 0

    def process_directory(self, input_dir, output_dir):
        """
        处理整个目录
        :param input_dir: 输入目录
        :param output_dir: 输出目录
        """
        os.makedirs(output_dir, exist_ok=True)
        file_paths = [
            os.path.join(input_dir, f)
            for f in os.listdir(input_dir)
            if f.endswith('.jsonl')
        ]

        total_tokens = 0
        for file_path in file_paths:
            tokens_count = self.process_file(file_path, output_dir)
            total_tokens += tokens_count

        logger.info(f"处理完成! 总token数: {total_tokens:,}")


    # 修改process_file方法支持并行
    def process_directory_parallel(self, input_dir, output_dir):
        from concurrent.futures import ProcessPoolExecutor

        os.makedirs(output_dir, exist_ok=True)
        file_paths = [os.path.join(input_dir, f) for f in os.listdir(input_dir)
                      if f.endswith('.jsonl')]

        with ProcessPoolExecutor(max_workers=self.config['num_processes']) as executor:
            results = list(tqdm(
                executor.map(self.process_file, file_paths, [output_dir] * len(file_paths)),
                total=len(file_paths)
            ))

        total_tokens = sum(r for r in results if r is not None)
        logger.info(f"处理完成! 总token数: {total_tokens:,}")

# 配置参数 (根据实际情况调整)
CONFIG = {
    # 输入输出
    'input_dir': 'path/to/your/jsonl_data',
    'output_dir': 'path/to/processed_npy',

    # Tokenizer设置
    'tokenizer_name': 'Qwen/Qwen3-4B-Base',
    'use_fast_tokenizer': True,
    'max_seq_length': 4096,  # 最大序列长度

    # 清洗参数
    'min_chars': 100,  # 最小字符数
    'max_chars': 10000,  # 最大字符数
    'min_chinese_ratio': 0.7,  # 最小中文比例
    'max_symbol_ratio': 0.3,  # 最大符号比例
    'max_repeat_ratio': 0.5,  # 最大重复内容比例

    # 安全过滤
    'sensitive_words_path': 'path/to/sensitive_words.txt',  # 敏感词表路径

    # 并行处理 (可选)
    'num_processes': multiprocessing.cpu_count(),
}

if __name__ == "__main__":
    preprocessor = ChineseCorpusPreprocessor(CONFIG)
    preprocessor.process_directory(CONFIG['input_dir'], CONFIG['output_dir'])

    # #test
    # # 加载tokenizer
    # tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Base", trust_remote_code=True)
    #
    # # 加载npy文件
    # tokens = np.load("path/to/processed_npy/file.npy")
    #
    # # 解码前100个token
    # decoded_text = tokenizer.decode(tokens[:100])
    # print(decoded_text)
    #
    # # 检查token分布
    # unique, counts = np.unique(tokens, return_counts=True)
    # print(f"总token数: {len(tokens):,}")
    # print(f"唯一token数: {len(unique):,}")